#!/usr/bin/env python
# -*-coding:utf-8 -*-

__all__ = ['Scheduler', '_determineFeasiblePath']

import logging
import time
import numpy as np
from task import Task
from time_window import TimeWindow
from robot_controller import RobotController
from linked_list import LinkedList


class Scheduler(object):
    """
    调度器。
    参考文献：
    Smolic-Rocak N, Bogdan S, Kovacic Z, et al. Time Windows Based Dynamic Routing in Multi-AGV Systems[J]. 
    IEEE Transactions on Automation Science & Engineering, 2009, 7(1):151-155.
    """
    def __init__(self, map, robots_manager, timer):
        """
        创建一个调度器。
        :param map: 机器人运行的地图。
        :param robots_manager: 机器人管理器。
        :param timer: 定时器。用于获取时间。
        """
        self.initLogger()
        self._map = map
        self._robots_manager = robots_manager
        self._timer = timer
        # 时间冗余量。同一边上，若两车前后脚进入的时间小于epsilon乘以在
        # 边上运动的时间，则认为两车之间的冲突。
        self._epsilon = 0.80
        #print('Scheduler timer:', self._timer)
        edges = self._map.getAllEdges()
        self._tw = TimeWindow(map, robots_manager.getRobotControllersNum())
        self._tw_backup = self._tw.copy()
        self._verbose = False
        
    def initLogger(self):
        self._logger = logging.getLogger('MyTCS.Scheduler')
    
    def getTimeWindow(self):
        return self._tw
        
    def backup(self):
        self._tw.copyto(self._tw_backup)
        
    def restore(self):
        self._tw_backup.copyto(self._tw)
    
    def _remainingTime(self, m, edge):
        """
        计算机器人m在边edge上的剩余时间（以最快速度移动所需时间）。
        :param m: 机器人的索引。
        :param edge: 连边。
        :return: 机器人m在边edge上的剩余时间。
        """
        controller = self._robots_manager.getRobotController(m)
        remaining_time = controller.getRemainingTimeOnCurrentEdge()
        if remaining_time == float('inf'):
            remaining_time = self._map.getEdgeWeight(edge)/controller.getMaxSpeed()
        return remaining_time
        
    def _nominalTraverseTime(self, m, edge):
        """
        计算机器人m在边edge上的标称时间（以最快速度移动所需时间）。
        :param m: 机器人的索引。
        :param edge: 连边。
        :return: 机器人m在边edge上的标称时间。
        """
        return self._map.getEdgeWeight(edge)/self._robots_manager.getRobotController(m).getMaxSpeed()
        
    def _isEdgeInPath(self, edge, path):
        if path is None:
            return False
            
        try:
            idx = path.index(edge[0])
        except ValueError:
            return False
        ret = (idx+1<len(path) and path[idx+1]==edge[1]) or (idx>0 and path[idx-1]==edge[1])
        if self._verbose:
            self._logger.debug('Edge {0} in path {1}: {2}'.format(edge, path, ret))
        return ret
      
    def _initTimeWindows(self, task, sigma, m, tm, robot2index):
        """
        初始化时间窗。对应于参考文献中的步骤一。
        :param task: 任务。Task对象。
        :param sigma: 路径，用路径上所有节点的列表表示。
        :param m: 机器人的索引（在机器人列表中索引）。
        :param tm: 当前时间。
        :param robot2index: 用于从机器人索引获取机器人上的任务的索引（按照优先级排序后的）。
                            比如robot2index={0:3, 1:2, 2:0, 3:1}表示机器人0上任务按优先级排序为3，机器人1排序为2，机器人2排序为0，优先级最高，机器人3排序为1。
        :return: 如果成功初始化，返回True，否则返回False。失败的情况为，当前任务的优先级高于了另一任务t，但执行t的机器人正好占用了当前任务需要使用的路段sigma上的边。
        """
        num = self._robots_manager.getRobotControllersNum()
        if self._verbose:
            self._logger.debug('STEP 1, initialization, time={0}'.format(tm))
            self._logger.debug('Time Window')
            self._tw.printForOccupiedEdges()
        
        for ii in range(num):
            #print('STEP 1: Before initialization, Exiting time of {0} is {1}'.format(aj, self._tw.t_o[aj]))
            #print('STEP 1: Current time is {0}'.format(tm))
            if ii == m:
                continue
                    
            controller = self._robots_manager.getRobotController(ii)
            ctsk = controller.getCompositeTask()

            for j in range(len(sigma)-1):
                aj = self._tw.disambiguate((sigma[j], sigma[j+1]))
                
                if ctsk and (not ctsk.isEmpty()):
                    tsk = ctsk.getNextUncompletedTask()
                    if self._tw.t_o[aj][ii] <= tm or \
                       (self._tw.t_i[aj][ii] < float('inf') and self._tw.t_i[aj][ii] > tm and \
                        ii in robot2index and \
                        robot2index[ii] > robot2index[m]):
                        self._clearTimeWindowForEdge(aj, ii)
                        if self._verbose and self._tw.t_o[aj][ii] <= tm:
                            self._logger.debug('STEP 1: clear time windows of robot {0} on edge {1} because its releasing time is before the current time'.format(ii, aj))
                        elif self._verbose and robot2index[ii] > robot2index[m]:
                            self._logger.debug('STEP 1: clear time windows of robot {0} on edge {1} because its task\'s priority {2} is lower than the current task {3}'.format( 
                                               ii, aj, tsk.getPriority(), task.getPriority()))
                                
            #print('STEP 1: {0}\'s entering time initialized to {1}'.format(aj, self._tw.t_i[aj]))
            #print('STEP 1: {0}\'s exiting time initialized to {1}'.format(aj, self._tw.t_o[aj]))
            #self._tw.w[aj][m] = self._remainingTime(m, (sigma[j], sigma[j+1]))
            #self._tw.t_i[aj][m] = float('inf')
            #self._tw.t_o[aj][m] = float('inf')
        om = self._tw.disambiguate((sigma[0], sigma[1]))
        self._tw.t_i[om][m] = tm
        self._tw.w[om][m] = self._remainingTime(m, (sigma[0], sigma[1]))
        self._tw.t_o[om][m] = tm + self._tw.w[om][m]
        if self._verbose:
            self._logger.debug('STEP 1: init robot {0}\'s entering and exiting time of {1}: {2}, {3}'.format(m, om, self._tw.t_i[om][m], self._tw.t_o[om][m]))
            self._logger.debug('STEP 1, initialization, time={0}'.format(tm))
            self._logger.debug('Time Window, initialized for robot {0} and path {1}'.format(m, sigma))
            self._tw.printForOccupiedEdges()
        return True
            
    def _insertTimeWindows(self, task, sigma, m, tm, a=1, q=0):
        """
        为边(sigma[a], sigma[a+1])插入时间窗。对应于参考文献中的步骤二。
        :param task: 任务。Task对象。
        :param sigma: 路径，用路径上所有节点的列表表示。
        :param m: 机器人的索引（在机器人列表中索引）。
        :param tm: 当前时间。
        :param a: 我们想为之插入时间窗的边在路径sigma中的索引。
        :param q: q是冲突边上，发生冲突的时间段的索引+1（详见_overlapTest的返回值说明）。
        """
        num = self._robots_manager.getRobotControllersNum()
        if a<=0 or q < 0:
            return False

        if self._verbose:
            self._logger.debug('STEP 2: current time={0}, tm={1}'.format(self._timer.getTime(), tm))
            self._logger.debug('STEP 2: inserting time windows for path {0}'.format(sigma))
        for j in range(a, len(sigma)-1):
            ai = self._tw.disambiguate((sigma[j-1], sigma[j]))
            aj = self._tw.disambiguate((sigma[j], sigma[j+1]))
            t_i_sorted = sorted(self._tw.t_i[aj])
            t_o_sorted = sorted(self._tw.t_o[aj])
            wmj_hat = self._nominalTraverseTime(m, (sigma[j], sigma[j+1]))
            if ((j == a and q == 0) or j != a) and \
               t_i_sorted[0] - (1.0+self._epsilon)*wmj_hat - self._tw.t_o[ai][m] > -1e-8:
                self._tw.t_i[aj][m] = self._tw.t_o[ai][m]
                self._tw.t_o[aj][m] = self._tw.t_i[aj][m] + wmj_hat
                if self._verbose:
                    self._logger.debug('STEP 2a: inserting robot {0}\'s entering and exiting time of {1}: {2}, {3}'.format(m, aj, self._tw.t_i[aj][m], self._tw.t_o[aj][m]))
            else:
                p = -1
                n = -1
                if t_i_sorted[0] != float('inf'):
                    for jj in range(num):
                        if jj == num - 1 or t_i_sorted[jj+1] == float('inf'):
                            n = jj
                            break
                            
                start = 0
                if j == a:
                    start = q
                    
                for jj in range(start, n+1):
                    if t_i_sorted[jj] - max(t_o_sorted[jj-1], self._tw.t_o[ai][m]) - (1+2*self._epsilon)*wmj_hat > -1e-8:
                        p = jj
                        break

                if p > 0:
                    self._tw.t_i[aj][m] = max(t_o_sorted[p-1]+self._epsilon*wmj_hat, self._tw.t_o[ai][m])
                    self._tw.t_o[aj][m] = self._tw.t_i[aj][m] + wmj_hat
                    if self._verbose:
                        self._logger.debug('STEP 2b: inserting robot {0}\'s entering and exiting time of {1}: {2}, {3}'.format(m, aj, self._tw.t_i[aj][m], self._tw.t_o[aj][m]))
                else:
                    if n >= 0:
                        self._tw.t_i[aj][m] = max(t_o_sorted[n] + self._epsilon*wmj_hat, self._tw.t_o[ai][m])
                    else:
                        self._tw.t_i[aj][m] = self._tw.t_o[ai][m]
                    self._tw.t_o[aj][m] = self._tw.t_i[aj][m] + wmj_hat
                    if self._verbose:
                        self._logger.debug('STEP 2c: inserting robot {0}\'s entering and exiting time of {1}: {2}, {3}'.format(m, aj, self._tw.t_i[aj][m], self._tw.t_o[aj][m]))
        return True
  
    def _clearTimeWindowForEdge(self, edge, m):
        """
        清除时间窗中，机器人m在边edge上的占用。
        :param edge: 边。
        :param m: 机器人的索引（在机器人列表中索引）。
        """
        e = self._tw.disambiguate(edge)
        self._tw.t_i[e][m] = float('inf')
        self._tw.t_o[e][m] = float('inf')
        self._tw.w[e][m] = 0.0
        
    def _clearTimeWindowForRobot(self, robot_controller):
        """
        清除时间窗中，机器人m在所有边占用的时间。
        :param edge: 边。
        :param m: 机器人的索引（在机器人列表中索引）。
        """
        path = robot_controller.getScheduledPath()
        m = robot_controller.getControlledRobotID()
        if path is not None:
            self._clearTimeWindowForPath(path, m)
    
    def _clearTimeWindowForPath(self, path, m):
        """
        清除时间窗中，机器人m在路径path上所有边的占用。
        :param path: 路径。
        :param m: 机器人的索引（在机器人列表中索引）。
        """
        for i in range(len(path)-1):
            e = (path[i], path[i+1])
            self._clearTimeWindowForEdge(e, m)

    def _overlapTest(self, task, sigma, m):
        """
        检测在插入时间窗之后，是否有冲突发生。
        :param task: 任务。Task对象。
        :param sigma: 路径，用路径上所有节点的列表表示。
        :param m: 机器人的索引（在机器人列表中索引）。
        :return: (j, q)，j是sigma中发生了冲突的边的索引。q是冲突边上，发生冲突的时间段的索引+1。
        """
        num = self._robots_manager.getRobotControllersNum()
        #print('STEP 3: sigma={0}, m={1}'.format(sigma, m))
        for j in range(len(sigma)-1):
            aj = self._tw.disambiguate((sigma[j], sigma[j+1]))
            if j < len(sigma)-2: # if aj is not the last path in sigma
                ai = self._tw.disambiguate((sigma[j+1], sigma[j+2]))
                # elongation
                if j == 0: # if aj is the first path in sigma
                    wmj_hat = self._remainingTime(m, (sigma[j], sigma[j+1]))
                else:
                    wmj_hat = self._nominalTraverseTime(m, (sigma[j], sigma[j+1]))
                wmj = wmj_hat + self._tw.t_i[ai][m] - self._tw.t_o[aj][m]
            else:
                if j == 0: # if aj is both the first and the last path in sigma
                    wmj = self._remainingTime(m, (sigma[j], sigma[j+1]))
                else:
                    wmj = self._nominalTraverseTime(m, (sigma[j], sigma[j+1]))
                    
            t_o_enlonged = np.copy(self._tw.t_o[aj])
            t_o_enlonged[m] = wmj + self._tw.t_i[aj][m]
            if self._verbose:
                self._logger.debug('STEP 3: exiting time enlonged to {0:.2f} on robot {1}, edge {2}'.format(t_o_enlonged[m], m, aj))

            # check overlaps
            t_i_sorted = sorted(self._tw.t_i[aj])
            t_o_enlonged_sorted = sorted(t_o_enlonged)

            n = num-1
            for jj in range(num-1):
                if t_i_sorted[jj+1] == float('inf'):
                    n = jj
                    break
            q = -1
            #print(t_i_sorted, n, t_o_enlonged_sorted)
            if n >= 0:
                for jj in range(n-1, -1, -1):
                    if t_i_sorted[jj+1]-t_o_enlonged_sorted[jj] < 0:
                        q = jj+1
                        self._clearTimeWindowForPath(sigma[j:], m)
                        if self._verbose:
                            self._logger.debug('STEP 3: confliction detected on robot {0}, edge {1}'.format(m, aj))
                        break
            if q != -1:
                if self._verbose:
                    self._logger.debug('STEP 3: confliction detected on edge {0}, reinsert after {1} time-windows on this edge'.format(aj, q))
                return j, q
        if self._verbose:
            self._logger.debug('STEP 3: confliction not detected')
        return None, None
        
    def schedule(self, task, robot_controller, paths, robot2index, verbose=False):
        """
        进行调度。
        :param task: 任务。Task对象。
        :param robot_controller: 任务所分配到的机器人的控制器。
        :param paths: 从任务task的源节点到目的节点的K条最短路径。
        :param robot2index: 用于从机器人索引获取机器人上的任务的索引（按照优先级排序后的）。
                            比如robot2index={0:3, 1:2, 2:0, 3:1}表示机器人0上任务按优先级排序为3，机器人1排序为2，机器人2排序为0，优先级最高，机器人3排序为1。
        :return: 路径，及在路径上到达各个节点的时间。
                 如果robot_controller为空，返回None, None;
                 如果paths中所有路径最后的驶出时间都为无穷大，返回None, None;
                 否则，返回paths中，驶出时间最早的那条路径，已经到达这条路径上每个点的时间。
        """
        self._verbose = verbose
        return self._determineFeasiblePath(self._timer.getTime(), task, robot_controller, paths, robot2index)

    def _determineFeasiblePath(self, tm, task, robot_controller, paths, robot2index):
        """
        确定一条可行路径，及在路径上到达各个节点的时间。对应于参考文献中的步骤三。
        :param tm: 当前时间。
        :param task: 任务。Task对象。
        :param robot_controller: 任务所分配到的机器人的控制器。
        :param paths: 从任务task的源节点到目的节点的K条最短路径。
        :param robot2index: 用于从机器人索引获取机器人上的任务的索引（按照优先级排序后的）。
                            比如robot2index={0:3, 1:2, 2:0, 3:1}表示机器人0上任务按优先级排序为3，机器人1排序为2，机器人2排序为0，优先级最高，机器人3排序为1。
        :return: 路径，及在路径上到达各个节点的时间。
                 如果robot_controller为空，返回None, None;
                 如果paths中所有路径最后的驶出时间都为无穷大，返回None, None;
                 否则，返回paths中，驶出时间最早的那条路径，已经到达这条路径上每个点的时间。
        """
        if robot_controller is None:
            return None, None
        m = robot_controller.getControlledRobotID()

        self._clearTimeWindowForRobot(robot_controller)
        waypoint_time = []
        w = [] # weight of paths
        for sigma in paths:
            #print('evaluating path {0}'.format(sigma))
            self._tw.backupForEdgesOnPath(sigma)
            ## STEP 1: Initialization of Time Vectors
            success = self._initTimeWindows(task, sigma, m, tm, robot2index)
            if not success:
                w.append(float('inf'))
                waypoint_time.append([])
                continue

            a = 1
            q = 0
            while True:
                ## STEP 2: Insertion of Time Windows
                self._insertTimeWindows(task, sigma, m, tm, a, q)
                
                ## STEP 3: Time Windows Elongation and Overlaps
                a, q = self._overlapTest(task, sigma, m)
                if a == 0:
                    # confliction detected on the first arc
                    w.append(float('inf'))
                    waypoint_time.append(None)
                    if self._verbose:
                        self._logger.debug('robot {0}, confliction detected on the 1st path!'.format(m))
                    break
                if a is None and q is None:
                    w.append(self._tw.t_o[self._tw.disambiguate((sigma[-2], sigma[-1]))][m])
                    waypoint_time.append([])
                    for j in range(len(sigma)-1):
                        aj = self._tw.disambiguate((sigma[j], sigma[j+1]))
                        waypoint_time[-1].append(self._tw.t_i[aj][m])
                    waypoint_time[-1].append(w[-1])
                    break
            self._tw.restoreForEdgesOnPath(sigma)
        #print(paths, w)
        min_w = float('inf')
        j = -1
        for i in range(len(paths)):
            if w[i] < min_w:
                j = i
                min_w = w[i]
        if j == -1:
            return None, None
        else:
            self._initTimeWindows(task, paths[j], m, tm, robot2index)
            for i in range(len(paths[j])-1):
                ai = self._tw.disambiguate((paths[j][i], paths[j][i+1]))
                self._tw.t_i[ai][m] = waypoint_time[j][i]
                self._tw.t_o[ai][m] = waypoint_time[j][i+1]
                if self._verbose:
                    self._logger.debug('Post-schedule: entering time {0:.2f}'.format(self._tw.t_i[ai][m]))
                    self._logger.debug('Post-schedule: exiting time {0:.2f}'.format(self._tw.t_o[ai][m]))
                self._tw.w[ai][m] = self._tw.t_o[ai][m] - self._tw.t_i[ai][m]
            return paths[j], waypoint_time[j]
